import os
import re
import sys
import cv2 as cv
import numpy as np
from PyCmpltrtok.common import rand_color, sep
from PyCmpltrtok.common_opencv import imzoom_rect

NUM = 1
TGT_RECT = (1500 // NUM // 2, 700)

image_path_file = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/data/voc.2012_mini.test'
result_dir = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/results/voc.2012_mini.test'
label_dir = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/scripts/VOCdevkit/VOC2012/labels'

# category names
label_names_path = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/data/voc.names'
with open(label_names_path, 'r') as f:
    names = f.readlines()
names = [name[:-2] if name[-2] == '\r' else name[:-1] for name in names]

# get all img paths
img_paths = []
with open(image_path_file, 'r', encoding='utf8') as f:
    while True:
        xpath = f.readline()
        if not xpath:
            break
        img_paths.append(xpath)
# print(img_paths)
for i, img_path in enumerate(img_paths):
    if img_path[-2:] == '\r\n':
        img_paths[i] = img_path[:-2]
        continue
    elif img_path[-1:] == '\n':
        img_paths[i] = img_path[:-1]
        continue
# print(img_paths)

# map main name to path and to boxes
name_list = []
name2path, name2boxes = {}, {}
for img_path in img_paths:
    dir_name, base_name = os.path.split(img_path)
    main_name, ext_name = os.path.splitext(base_name)
    name_list.append(main_name)
    name2path[main_name] = img_path
    name2boxes[main_name] = []

# read in box info
result_file_names = os.listdir(result_dir)
regexp = re.compile(r'_([^\._]+)\.txt$')
for result_file_name in result_file_names:
    result_file_path = os.path.join(result_dir, result_file_name)
    matcher = regexp.search(result_file_name)
    if not matcher:
        print(f'File {result_file_name} named badly!', file=sys.stderr)
        continue
    label = matcher.group(1)
    boxes = []
    with open(result_file_path, 'r') as f:
        while True:
            line = f.readline()
            if not line:
                break
            vals = line.split()
            main_name = vals[0]
            objness = vals[1]
            box_info = [label]
            box_info.extend([int(np.round(float(x))) for x in vals[2:]])
            box_info.append(float(objness))
            name2boxes[main_name].append(box_info)


# read in box info
def get_boxes(xid):
    label_path = os.path.join(label_dir, f'{xid}.txt')
    with open(label_path, 'r') as f:
        lines = f.readlines()
    boxes = []
    lines = [line[:-2] if line[-2] == '\r' else line[:-1] for line in lines]
    for line in lines:
        vals = line.split()
        name = names[int(vals[0])]
        x, y, w, h = [float(val) for val in vals[1:]]
        box_info = [name]
        box_info.extend([x - w/2, y - h/2, x + w/2, y + h/2])
        boxes.append(box_info)
    return boxes


# calc IOU
def calc_iou(x1, y1, x2, y2, xlist):
    iou_max = 0
    for x1ref, y1ref, x2ref, y2ref in xlist:
        max_x1 = max(x1, x1ref)
        max_y1 = max(y1, y1ref)
        min_x2 = min(x2, x2ref)
        min_y2 = min(y2, y2ref)
        intersected_x = max(0, min_x2 - max_x1)
        intersected_y = max(0, min_y2 - max_y1)
        intersected = intersected_x * intersected_y
        iou = intersected / ((x2 - x1) * (y2 - y1) + (x2ref - x1ref) * (y2ref - y1ref) - intersected)
        if iou > iou_max:
            iou_max = iou
    return iou_max


xpos = 0
xlen = len(name_list)
while True:
    xnames = name_list[xpos:xpos+NUM]
    img_arr = []
    for xid in xnames:
        sep(f'#{xpos} {xid}')
        xpath = name2path[xid]
        img = cv.imread(xpath, cv.IMREAD_COLOR)
        img_ = img.copy()
        pred_label2box, gt_label2box = {}, {}

        # predictions
        xboxes_pred = name2boxes[xid]
        for box in xboxes_pred:
            label, x1, y1, x2, y2, objness = box

            # prepare data to calc IOU
            xlist = pred_label2box.get(label, [])
            xlist.append([x1, y1, x2, y2])
            pred_label2box[label] = xlist

            color = rand_color()
            cv.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv.putText(img, f'{label} ({objness:.4f})', (x1, y1 + 15), cv.FONT_HERSHEY_PLAIN, 1, color, 1)
        img = imzoom_rect(img, TGT_RECT)
        color = rand_color()
        cv.putText(img, f'#{xpos} {xid} (predictions)', (0, 15), cv.FONT_HERSHEY_PLAIN, 1, color, 2)
        img_arr.append(img)

        # ground truth
        img = img_.copy()
        xboxes = get_boxes(xid)
        h, w = img.shape[:2]
        xboxes_gt = []
        for box in xboxes:
            label = box[0]
            xy12 = np.array(box[1:]) * np.array([w, h, w, h])
            xy12 = np.round(xy12).astype(int)
            x1, y1, x2, y2 = xy12
            xboxes_gt.append([label, x1, y1, x2, y2])

            # prepare data to calc IOU
            xlist = gt_label2box.get(label, [])
            xlist.append([x1, y1, x2, y2])
            gt_label2box[label] = xlist

            color = rand_color()
            cv.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv.putText(img, f'{label}', (x1, y1 + 15), cv.FONT_HERSHEY_PLAIN, 1, color, 2)
        img = imzoom_rect(img, TGT_RECT)
        color = rand_color()
        cv.putText(img, f'#{xpos} {xid} (GT)', (0, 15), cv.FONT_HERSHEY_PLAIN, 1, color, 2)
        img_arr.append(img)

        # print info with IOU
        iou_pred = []
        for i, (label, x1, y1, x2, y2, objness) in enumerate(xboxes_pred):
            iou = calc_iou(x1, y1, x2, y2, gt_label2box.get(label, []))
            iou_pred.append(iou)
        iou_gt = []
        for i, (label, x1, y1, x2, y2) in enumerate(xboxes_gt):
            iou = calc_iou(x1, y1, x2, y2, pred_label2box.get(label, []))
            iou_gt.append(iou)
        for i, (label, x1, y1, x2, y2, objness) in enumerate(xboxes_pred):
            print(f'{label} {x1, y1, x2, y2}, obj={objness:.4f}, IOU={iou_pred[i]:.4f}')
        print('GT:')
        for i, (label, x1, y1, x2, y2) in enumerate(xboxes_gt):
            print(f'{label} {x1, y1, x2, y2}, IOU={iou_gt[i]:.4f}')

    img = np.concatenate(img_arr, axis=1)
    cv.line(img, (TGT_RECT[0], 1), TGT_RECT, (0, 255, 0), 2)
    cv.imshow('check YOLO valid', img)
    k = cv.waitKey()
    if ord('p') == k:
        xpos -= 1
    elif ord('n') == k:
        xpos += 1
    elif ord('q') == k:
        print('Bye!')
        break
    xpos = xpos if xpos >= 0 else 0
    xpos = xpos if xpos < xlen else xlen - 1
    cv.destroyAllWindows()
